// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright @ Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: kernel hot upgrade
 */
#include <linux/types.h>
#include <linux/init.h>
#include <linux/export.h>
#include <linux/kstrtox.h>
#include <linux/kthread.h>
#include <linux/jiffies.h>
#include <linux/module.h>
#include <linux/kallsyms.h>

u64 kernel_hotupgrade;
EXPORT_SYMBOL(kernel_hotupgrade);

static int __init set_kernel_hotupgrade(char *str)
{
	int err;

	if (!str)
		return 0;

	err = kstrtoull(str, 0, &kernel_hotupgrade);
	if (err < 0)
		return 0;

	return 1;
}

__setup("kernel_hotupgrade=", set_kernel_hotupgrade);

void __ref free_init_mem_async(void)
{
	struct task_struct *kthread;

	if (!(HOTUPGRADE_WITH(UPGRADE_F_DELAY_INITCALL))) {
		if (HOTUPGRADE_WITH(UPGRADE_F_ASYNC_FREE_INITMEM)) {
			kthread = kthread_run(free_init_mem, NULL, "async_free_initmem");
			if (IS_ERR(kthread)) {
				pr_warn("cannot start thread to free initmem, free initmem directly\n");
				free_init_mem(NULL);
			}
		} else {
			free_init_mem(NULL);
		}
	}
}

/*
 * The initcall should keep the original called order.
 */
static struct delay_initcall_info delayed_initcalls[] = {
	{
		.name = "init_acpi_pm_clocksource",
	},
	{
		.name = "efifb_driver_init",
	},
	{
		.name = "serial8250_init",
	},
	{
		.name = "cacheinfo_sysfs_init",
	},
	{
		.name = "mcheck_init_device",
	},
	{
		.name = "resctrl_late_init",
	},
	{
		.name = "init_encrypted",
	},
	{
		.name = "load_uefi_certs",
	},
	{
		.name = "late_free_initrd",
	},
};

static int __init_or_module initcall_delaylisted(initcall_t fn, unsigned long len)
{
	char fn_name[KSYM_SYMBOL_LEN];
	unsigned long addr;
	unsigned long i;

	if (!len)
		return -1;

	addr = (unsigned long) dereference_function_descriptor(fn);
	sprint_symbol_no_offset(fn_name, addr);
	strreplace(fn_name, ' ', '\0');
	for (i = 0; i < len; ++i) {
		if (!strcmp(fn_name, delayed_initcalls[i].name)) {
			pr_debug("initcall %s delaylisted\n", fn_name);
			return i;
		}
	}

	return -1;
}

bool __init set_delayed_initcall(initcall_t fn)
{
	bool ret = false;
	int index;
	unsigned long len = sizeof(delayed_initcalls) / sizeof(*delayed_initcalls);

	if (HOTUPGRADE_WITH(UPGRADE_F_DELAY_INITCALL)) {
		index = initcall_delaylisted(fn, len);
		if (index >= 0) {
			delayed_initcalls[index].call = fn;
			ret = true;
		}
	}

	return ret;
}

static void do_delay_initcalls(void)
{
	unsigned long i, len;

	len = sizeof(delayed_initcalls) / sizeof(*delayed_initcalls);
	if (!len)
		return;

	for (i = 0; i < len; ++i) {
		if (delayed_initcalls[i].call)
			do_one_initcall(delayed_initcalls[i].call);
	}
}

#define SCHEDULE_DELAYED_TIME	100
bool initcall_delaying;
static void do_last_initcalls(struct work_struct *work);
static DECLARE_DELAYED_WORK(last_init_work, do_last_initcalls);
static void do_last_initcalls(struct work_struct *work)
{
	if (kernel_hotupgrade == 0) {
		do_delay_initcalls();
		free_init_mem(NULL);
		initcall_delaying = false;
	} else {
		schedule_delayed_work(&last_init_work, msecs_to_jiffies(SCHEDULE_DELAYED_TIME));
	}
}

void __init schedule_delayed_initcalls(void)
{
	if (HOTUPGRADE_WITH(UPGRADE_F_DELAY_INITCALL)) {
		initcall_delaying = true;
		schedule_delayed_work(&last_init_work, msecs_to_jiffies(SCHEDULE_DELAYED_TIME));
	}
}
